package com.tianji.aigc.service.impl;

import cn.hutool.core.collection.CollUtil;
import com.tianji.aigc.dto.ChatDTO;
import com.tianji.aigc.entity.ChatSession;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.memory.RedisChatMemory;
import com.tianji.aigc.service.IChatService;
import com.tianji.aigc.service.IChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.exceptions.BadRequestException;
import com.tianji.common.utils.UserContext;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class ChatServiceImpl implements IChatService {
    @Autowired
    private ChatClient chatClient;

    @Autowired
    private IChatSessionService chatSessionService;

    /**
     * 定义可以控制连续输出标识值的Map，Map的KEY是sessionId，Vlue就是布尔值，true表示可以继续输出，false表示中断输出
     */
    private static Map<String,Boolean> CONTINUE_FLAG_MAP = new ConcurrentHashMap<>();

    /**
     * 创建流式对话
     *
     * @param chatDTO ;
     */
    @Override
    public Flux<ChatEventVO> chat(ChatDTO chatDTO) {
        //首先判断会话是否存在
        String seesionId = chatDTO.getSessionId();
        Long count = chatSessionService.lambdaQuery().eq(ChatSession::getSessionId, chatDTO.getSessionId()).count();
        if (count==0){
            throw new BadRequestException("请先创建会话");
        }
        //拼接对话id，由用户id和会话id拼接成
        String conversationId = UserContext.getUser() +":"+seesionId;
        StringBuffer content =new StringBuffer();//用于流式输出，StringBuffer拼接时候是线程安全的
        //进行模型会话
        return chatClient.prompt()
                .user(chatDTO.getQuestion())//获取用户的问题
                .advisors(x->x.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,conversationId))
                .stream()
                .content()
                .doFirst(()->{CONTINUE_FLAG_MAP.put(seesionId,true);})//标记开始输出第一步，设置为true，能够继续输出
                .doOnError(x->CONTINUE_FLAG_MAP.remove(seesionId))//当生成错误的时候就会删除map，节省空间
                .doOnComplete(()->{CONTINUE_FLAG_MAP.remove(seesionId);})//输出完成，删除map节省空间
                //接下来是异常的情况
                .doOnCancel(()->{this.saveToRedis(conversationId,content.toString());})//当发现取消输出时，就将之前收集到的信息保存到redis当中
                .takeWhile(x->CONTINUE_FLAG_MAP.getOrDefault(seesionId,false))//根据布尔值来判断是否继续输出，true表示继续，false表示终止
                .map(x->{
                    content.append(x);
                    return ChatEventVO.builder()
                            .eventData(x)
                            .eventType(ChatEventTypeEnum.DATA.getValue())
                            .build();
                })
                .concatWith(Flux.just(ChatEventVO.builder()
                                .eventType(ChatEventTypeEnum.STOP.getValue())//结束输出
                        .build()));


    }
    @Autowired
    private RedisChatMemory redisChatMemory;
    private void saveToRedis(String conversationId, String content) {
        redisChatMemory.add(conversationId,new AssistantMessage(content));
    }

    /**
     * 停止会话
     *
     * @param sessionId ;
     */
    @Override
    public void stop(String sessionId) {
        //核心思路：打断stream流式输出，不让继续输出，必须改变可以连续输出的布尔标识的值，将值设置为false
        CONTINUE_FLAG_MAP.put(sessionId,false);
    }
}
